package main

import (
	"bytes"
	"context"
	"encoding/hex"
	"errors"
	"flag"
	"gitee.com/leaf-rain-gitee/util/tool"
	"github.com/ethereum/go-ethereum/accounts/keystore"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/panjf2000/ants/v2"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/gridfs"
	"go.mongodb.org/mongo-driver/mongo/options"
	"go.mongodb.org/mongo-driver/mongo/readpref"
	"go.uber.org/atomic"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"log"
	"sync"
	"time"
)

var dbClientAcc *gorm.DB
var dbClientPri *gorm.DB
var mgoClient *mongo.Client
var count atomic.Int64

//var path = flag.String("path", "./pwd.txt", "path")
var limit = flag.Int("limit", 1040633, "limit")
var offset = flag.Int("offset", 0, "offset")
var dsnAcc = flag.String("dsnAcc", "root:123456@tcp(172.12.17.160:3306)/gluttonous?timeout=5s&readTimeout=5s&writeTimeout=5s&parseTime=true&loc=Local&charset=utf8,utf8mb4", "dsnAcc")
var dsnPri = flag.String("dsnPri", "root:+eN(2dFc5qu.@tcp(172.12.17.165:3306)/gluttonous?timeout=5s&readTimeout=5s&writeTimeout=5s&parseTime=true&loc=Local&charset=utf8,utf8mb4", "dsnPri")
var mgo = flag.String("mgo", "mongodb://172.12.17.78:27017", "mgo")

func init() {
	var err error
	dbClientAcc, err = gorm.Open(mysql.Open(*dsnAcc))
	if err != nil {
		panic(err)
	}
	dbClientPri, err = gorm.Open(mysql.Open(*dsnPri))
	if err != nil {
		panic(err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	mgoClient, err = mongo.Connect(ctx, options.Client().ApplyURI(*mgo))
	ctx, cancel = context.WithTimeout(context.Background(), 120*time.Second)
	defer cancel()
	err = mgoClient.Ping(ctx, readpref.Primary())
	if err != nil {
		panic(err)
	}
}

type Param struct {
	Addr    string `json:"addr"`
	Private string `json:"private"`
}

func myFunc(param interface{}) {
	var err error
	var addr = param.(string)
	var private string

	bucket, _ := gridfs.NewBucket(mgoClient.Database("keystore"), options.GridFSBucket().SetName("file"))
	fileBuffer := bytes.NewBuffer(nil)
	if _, err = bucket.DownloadToStream(addr[2:], fileBuffer); err != nil && !errors.Is(err, gridfs.ErrFileNotFound) {
		panic(err)
	}
	data := fileBuffer.Bytes()
	private, err = NewTransferOptByKeyJson(data, "HX141656")
	if err != nil {
		panic(err)
	}

	err = dbClientPri.Exec("insert into gs_wallet_private(wallet_address, private_key, create_time)values (?, ?, ?)", addr, private, tool.GetTimeUnixMilli()).Error
	if err != nil {
		log.Println("insert failed:", addr, ":", private)
	} else {
		log.Println("insert success:", addr, ":", private)
	}
}

func main() {
	flag.Parse()
	var ctx = context.Background()
	if *limit == 0 {
		panic(errors.New("limit zero error"))
	}
	defer ants.Release()
	var wg sync.WaitGroup
	p, _ := ants.NewPoolWithFunc(4, func(i interface{}) {
		myFunc(i)
		wg.Done()
	})
	defer p.Release()
	defer func() {
		if err := mgoClient.Disconnect(ctx); err != nil {
			panic(err)
		}
	}()

	// 从数据库中读取数据
	var wallets = []Wallet{}
	err := dbClientAcc.Raw("select `gs_account`.`wallet` from `gs_account` where 1=1  limit ? offset ?", *limit, *offset).Find(&wallets).Error
	if err != nil {
		log.Println("db exec select wallet failed:", *limit, *offset)
		panic(err)
	}
	var walletAddr string
	for i := range wallets {
		walletAddr = wallets[i].Wallet
		wg.Add(1)
		_ = p.Invoke(walletAddr)
	}
	wg.Wait()
}

func NewTransferOptByKeyJson(keyJson []byte, passWard string) (string, error) {
	key, err := keystore.DecryptKey(keyJson, passWard)
	if err != nil {
		return "", err
	}
	return hex.EncodeToString(crypto.FromECDSA(key.PrivateKey)), nil
}

type Item struct {
	ID            int    `json:"id" gorm:"primaryKey;column:id;type:int auto_increment;comment:id"`                                                                       // id
	WalletAddress string `json:"wallet_address" gorm:"column:wallet_address;type:varchar(128);not null;uniqueIndex:gs_wallet_private_wallet_address_uindex;comment:钱包地址"` // 钱包地址
	PrivateKey    string `json:"private_key" gorm:"column:private_key;type:varchar(128);not null;comment:私钥"`                                                             // 私钥
	CreateTime    int64  `json:"create_time" gorm:"column:create_time;type:bigint;not null;default:0;comment:创建时间"`                                                       // 创建时间
}

// TableName returns the table name of the Item model
func (i *Item) TableName() string {
	return "gs_wallet_private"
}

type Wallet struct {
	Wallet string `json:"wallet" gorm:"column:wallet;type:varchar(100);not null;uniqueIndex:gs_account_wallet_uindex;comment:钱包地址"` // 钱包地址 	// 修改时间
}
